from transformers import M2M100ForConditionalGeneration, M2M100Tokenizer

hi_text = "जीवन एक चॉकलेट बॉक्स की तरह है।"
chinese_text = "生活就像一盒巧克力。"
en_txt = "I can help create an image inspired by this style. To replicate this look, could you describe any specific details you’d like to include, such as the characters’ appearance, background elements, or any particular emotions or theme you'd like emphasized?"

model = M2M100ForConditionalGeneration.from_pretrained("../m2m100_418M")
tokenizer = M2M100Tokenizer.from_pretrained("../m2m100_418M")

# # translate Hindi to French
# tokenizer.src_lang = "hi"
# encoded_hi = tokenizer(hi_text, return_tensors="pt")
# generated_tokens = model.generate(**encoded_hi, forced_bos_token_id=tokenizer.get_lang_id("fr"))
# decoded_sequences = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
# # => "La vie est comme une boîte de chocolat."
# print(decoded_sequences)
#
# # translate Chinese to English
# tokenizer.src_lang = "zh"
# encoded_zh = tokenizer(chinese_text, return_tensors="pt")
# generated_tokens = model.generate(**encoded_zh, forced_bos_token_id=tokenizer.get_lang_id("en"))
# decoded_sequences = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
# # => "Life is like a box of chocolate."
# print(decoded_sequences)

# translate English to Chinese
tokenizer.src_lang = "en"
encoded_en = tokenizer(en_txt, return_tensors="pt")
generated_tokens = model.generate(**encoded_en, forced_bos_token_id=tokenizer.get_lang_id("zh"))
decoded_sequences = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
# => "Life is like a box of chocolate."
print(decoded_sequences)